import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

cluster_count = 256

data_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'data'))
features_path = os.path.join(data_path, 'image_features')

datas = []
features = []
for i in range(100):
    path = os.path.join(features_path, f"cache_feature{i}.bin")
    with open(path, 'rb') as fp:
        tmp_datas = pickle.load(fp)
        for data in tmp_datas:
            datas.append(data)
            features.append(data['features'])

features = np.vstack(features)

print("start training")

kmeans = KMeans(n_clusters=cluster_count, n_init='auto', verbose=1)
kmeans.fit(features)

# print(kmeans.cluster_centers_)
# print(kmeans.cluster_centers_.shape)
# print(kmeans.labels_)
# print(kmeans.labels_.shape)
# print(kmeans.labels_.min())
# print(kmeans.labels_.max())

invert_indices = []
for i in range(cluster_count):
    cluster_info = { 'centroid' : kmeans.cluster_centers_[i], 'ids' : [] }
    invert_indices.append(cluster_info)
data_count = len(datas)
for i in range(data_count):
    label = kmeans.labels_[i]
    invert_indices[label]['ids'].append(datas[i]['id'])

with open(os.path.join(data_path, 'invert_indices.bin'), 'wb') as fp:
    pickle.dump(invert_indices, fp)

plt.hist(kmeans.labels_, bins=256)
plt.show()